#include "pch.h"
#include "beautycore.h"
#include "openglwrapper.h"

#include <opencv2/imgproc/imgproc.hpp>

CBeautyCore::CBeautyCore()
{

}
CBeautyCore::~CBeautyCore()
{

}

void CBeautyCore::RegisterCallback(RetriveFrameCallback callback)
{

}

cv::Mat CBeautyCore::ProcessFrame(const cv::Mat& frame)
{
	return opengl_wrapper_->ProcessFrame(frame);
}

bool CBeautyCore::ProcessFrame(uint8_t* y,uint8_t* u, uint8_t* v,
					int stride_y, int stride_u, int stride_v)
{
	std::shared_ptr<uint8_t[]> buff;
	buff.reset(new uint8_t[buff_size]);
	int offset = 0;

	// copy the data
	if(stride_y != width) {
		for(auto i = 0; i < height; i++) {
            memcpy((void*)(buff.get()[offset]), y + stride_y * i, width);
			offset += width;
		}
	} else {

        memcpy((void*)(buff.get()[offset]), y, width * height);
        offset += width * height;

    }

	if(stride_u != (width/2)) {
		for(auto i = 0; i < height/2; i++) {
			memcpy((void*)(buff.get()[offset]), u + stride_u * i, width/2);
			offset += width/2;
		}
	} else {
		memcpy((void*)(buff.get()[offset]), u, width * height /4);
		offset += width * height / 4;

	}

	if(stride_v != (width/2)) {
		for(auto i = 0; i < height/2; i++) {
			memcpy((void*)(buff.get()[offset]), v + stride_v * i, width/2);
			offset += width/2;
		}
	} else {
		memcpy((void*)(buff.get()[offset]), v, width * height /4);
		offset += width * height / 4;
	}

	cv::Mat yuv_frame(width, int(double(height) / 1.5), CV_8UC1, buff.get());
	cv::Mat bgr_frame;

	cv::cvtColor(yuv_frame, bgr_frame, cv::COLOR_YUV2BGR_I420);

	cv::Mat frame_after_proccess = opengl_wrapper_->ProcessFrame(bgr_frame);

	cv::Mat result_frame;
	cv::cvtColor(frame_after_proccess, result_frame, cv::COLOR_BGR2YUV_I420);

	// copy the data
	offset = 0;
	if(stride_y != width) {
		for(auto i = 0; i < height; i++) {
			memcpy( y + stride_y * i, (const void*)( result_frame.data[offset]), width);
			offset += width;
		}
	} else {
        memcpy(y, (const void*)(result_frame.data[offset]), width * height);
		offset += width * height;
	}

	if(stride_u != (width/2)) {
		for(auto i = 0; i < height/2; i++) {
			memcpy(u + stride_u * i, (const void*)(result_frame.data[offset]), width/2);
			offset += width/2;
		}
	} else {
		memcpy(u, (const void*)(result_frame.data[offset]), width * height /4);
		offset += width * height / 4;
	}

	if(stride_v != (width/2)) {
		for(auto i = 0; i < height/2; i++) {
			memcpy(v + stride_v * i,(const void*)( result_frame.data[offset]), width/2);
			offset += width/2;
		}
	} else {
		memcpy(v, (const void*)(result_frame.data[offset]), width * height /4);
		offset += width * height / 4;
	}

	return true;
}

void CBeautyCore::SetWhiteningParam(const float factor)
{
	if (opengl_wrapper_)
	{
		opengl_wrapper_->SetWhitenLevel(factor);
	}
}

bool CBeautyCore::InitializeBeautyCore(int w, int h)
{
	width = w;
	height = h;
	buff_size = int(double(width * height) * 1.5);

	opengl_wrapper_ = std::shared_ptr<COpenglWrapper>(new COpenglWrapper());
	return opengl_wrapper_->Initialize(w, h);
}
